import json
import argparse
import os
import open3d as o3d
import numpy as np
from my_io import mkdir
from tqdm import tqdm
from tictoc import TicToc
import multiprocessing


def main(args):
    # 统计耗时
    cost = TicToc("格式转换")
    assert os.path.exists(args.data_path)
    assert args.fun in ['pcd', 'json']
    datasets = os.listdir(args.data_path)
    datasets.sort(key=lambda x: x)
    files = []
    for dataset in tqdm(datasets):
        if "__" == dataset[:2]:  # 数据集以"__"开头
            dataset_path = os.path.join(args.data_path, dataset)
            files.append(dataset_path)

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(dataset_path, args):
    if args.fun == 'json':
        print("开始json转txt...")
        json_paths = os.path.join(dataset_path, "samples/label")
        txt_paths = os.path.join(dataset_path, "samples/label_txt")
        mkdir(txt_paths)

        for label in os.listdir(json_paths):
            filename = json_paths + '/' + label

            with open(filename, 'r') as f:
                data = json.load(f)

            label_list = []
            for obj_dict in data:
                label_name = obj_dict["obj_type"]
                pos_xyz = obj_dict["psr"]["position"]
                rot_xyz = obj_dict["psr"]["rotation"]
                scale_xyz = obj_dict["psr"]["scale"]

                temp = str(label_name) + " "
                for xyz_dict in [pos_xyz, scale_xyz]:
                    for key in ["x", "y", "z"]:
                        temp += str(xyz_dict[key])
                        temp += " "
                temp += str(rot_xyz["z"]) + "\n"
                label_list.append(temp)
            txt_name = txt_paths + '/' + str.split(label, '.')[0] + ".txt"
            with open(txt_name, "w") as f:
                for label in label_list:
                    f.write(label)

        print("<{}> json converted txt".format(dataset_path))

    elif args.fun == 'pcd':
        lidar_paths = os.path.join(dataset_path, "lidar")
        lidar_bin_paths = os.path.join(dataset_path, "lidar_bin")
        # radar_paths = os.path.join(dataset_path, "radar")
        # radar_bin_paths = os.path.join(dataset_path, "radar_bin")
        mkdir(lidar_bin_paths)
        # mkdir(radar_bin_paths)
        # for pcd_path in os.listdir(radar_paths):
        #     pcd = o3d.io.read_point_cloud(os.path.join(radar_paths, pcd_path))
        #     pc = np.asarray(pcd.points).astype(np.float32)
        #     pc.tofile(os.path.join(radar_bin_paths, str.split(pcd_path, '.')[0] + '.bin'))

        for pcd_path in os.listdir(lidar_paths):
            pcd = o3d.io.read_point_cloud(os.path.join(lidar_paths, pcd_path))
            pc = np.asarray(pcd.points).astype(np.float32)
            pc.tofile(os.path.join(lidar_bin_paths, str.split(pcd_path, '.')[0] + '.bin'))

        print("<{}> pcd converted bin".format(dataset_path))
    else:
        print('args.fun <{}> is error'.format(args.fun))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default='/media/adt/ZWH4T/ZWH/bags/dataset/tm3/work_space_tmp/datasets',
                        help='your data root path')
    parser.add_argument('--fun', default='json', help='pcd2bin or json2txt')
    args = parser.parse_args()

    main(args)
